#!/usr/bin/env python3
"""
FaceFusion FastAPI Server
简洁的人脸交换 API 服务
"""

import os
import uuid
import shutil
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# 设置环境变量
os.environ['OMP_NUM_THREADS'] = '1'

from facefusion import state_manager
from facefusion.args import apply_args
from facefusion.core import common_pre_check, processors_pre_check, conditional_process
from facefusion.filesystem import is_image, is_video
from facefusion.program import create_program
from facefusion.vision import detect_video_resolution, detect_video_fps, pack_resolution

# 创建 FastAPI 应用
app = FastAPI(title="FaceFusion API", version="1.0.0")

# CORS 支持
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置目录
UPLOAD_DIR = Path("./uploads")
OUTPUT_DIR = Path("./outputs")
UPLOAD_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)

# 任务状态
tasks = {}

# 全局初始化标志
_initialized = False


def init_facefusion():
    """初始化 FaceFusion 配置（模拟图形界面的行为）"""
    from facefusion.execution import has_execution_provider

    # 使用和图形界面完全相同的初始化方式
    program = create_program()

    # 模拟命令行参数，使用 headless-run 模式
    args = program.parse_args(['headless-run'])

    # 应用所有默认参数（和图形界面一样）
    apply_args(vars(args), state_manager.init_item)

    # 模拟用户在图形界面选择 CUDA（如果可用）
    if has_execution_provider('cuda'):
        state_manager.set_item('execution_providers', ['cuda'])
        print("✅ FaceFusion 配置初始化完成")
        print("   执行提供商: ['cuda']")
        print("   🚀 CUDA 加速已启用")
    else:
        state_manager.set_item('execution_providers', ['cpu'])
        print("✅ FaceFusion 配置初始化完成")
        print("   执行提供商: ['cpu']")
        print("   ⚠️  CUDA 不可用，使用 CPU")
        print("   💡 如需 GPU 加速，请运行: python install.py --onnxruntime cuda")

    # 确保关键的输出参数不为 None（headless-run 模式不会设置这些）
    # 注意：output_video_resolution 和 output_video_fps 会在处理时根据目标视频自动设置
    if state_manager.get_item('output_video_scale') is None:
        state_manager.set_item('output_video_scale', 100)
    if state_manager.get_item('output_video_quality') is None:
        state_manager.set_item('output_video_quality', 80)
    if state_manager.get_item('output_video_encoder') is None:
        state_manager.set_item('output_video_encoder', 'libx264')
    if state_manager.get_item('output_video_preset') is None:
        state_manager.set_item('output_video_preset', 'medium')


def process_video(source_img: Path, target_vid: Path, output_vid: Path, options: dict = None) -> tuple:
    """
    执行人脸交换

    Returns:
        (success: bool, message: str)
    """
    import traceback

    try:
        # 更新特定参数
        if options:
            for key, value in options.items():
                state_manager.set_item(key, value)

        # 设置路径
        state_manager.set_item('source_paths', [str(source_img)])
        state_manager.set_item('target_path', str(target_vid))
        state_manager.set_item('output_path', str(output_vid))

        # 自动获取目标视频的分辨率和帧率，使输出视频与目标视频保持一致
        video_resolution = detect_video_resolution(str(target_vid))
        video_fps = detect_video_fps(str(target_vid))

        if video_resolution:
            resolution_str = pack_resolution(video_resolution)
            state_manager.set_item('output_video_resolution', resolution_str)
            print(f"[INFO] 目标视频分辨率: {resolution_str}")

        if video_fps:
            state_manager.set_item('output_video_fps', video_fps)
            print(f"[INFO] 目标视频帧率: {video_fps} fps")

        # 预检查
        print("[DEBUG] 开始环境检查...")
        if not common_pre_check():
            return False, '环境检查失败：请确保已安装 ffmpeg 和 curl'

        print("[DEBUG] 开始处理器检查...")
        if not processors_pre_check():
            return False, '模型检查失败：请运行 python facefusion.py force-download'

        # 执行处理
        print("[DEBUG] 开始执行人脸交换...")
        error_code = conditional_process()

        if error_code == 0:
            return True, '处理完成'
        elif error_code == 2:
            return False, '未检测到人脸！请使用清晰的正面照片'
        elif error_code == 3:
            return False, '内容分析失败'
        else:
            return False, f'处理失败，错误码: {error_code}'

    except Exception as e:
        print(f"[ERROR] 异常: {str(e)}")
        print("[ERROR] 详细错误信息:")
        traceback.print_exc()
        return False, f'异常: {str(e)}'


@app.post("/api/face-swap")
async def face_swap(
    source_image: UploadFile = File(..., description="源人脸图片"),
    target_video: UploadFile = File(..., description="目标视频"),
    model: Optional[str] = Form("inswapper_128", description="模型"),
    quality: Optional[int] = Form(80, description="质量 0-100"),
    threads: Optional[int] = Form(4, description="线程数")
):
    """人脸交换 API"""
    task_id = f"task_{uuid.uuid4().hex[:12]}"

    try:
        # 保存文件
        src_path = UPLOAD_DIR / f"{task_id}_src{Path(source_image.filename).suffix}"
        tgt_path = UPLOAD_DIR / f"{task_id}_tgt{Path(target_video.filename).suffix}"
        out_path = OUTPUT_DIR / f"{task_id}_out.mp4"

        with open(src_path, "wb") as f:
            shutil.copyfileobj(source_image.file, f)

        with open(tgt_path, "wb") as f:
            shutil.copyfileobj(target_video.file, f)

        # 验证文件
        if not is_image(str(src_path)):
            raise HTTPException(400, "无效的图片文件")
        if not is_video(str(tgt_path)):
            raise HTTPException(400, "无效的视频文件")

        # 处理
        print(f"\n{'='*60}")
        print(f"任务: {task_id}")
        print(f"源图片: {source_image.filename}")
        print(f"目标视频: {target_video.filename}")
        print(f"{'='*60}")

        options = {
            'face_swapper_model': model,
            'output_video_quality': quality,
            'execution_thread_count': threads
        }

        success, message = process_video(src_path, tgt_path, out_path, options)

        # 清理上传文件
        src_path.unlink(missing_ok=True)
        tgt_path.unlink(missing_ok=True)

        if success and out_path.exists():
            tasks[task_id] = str(out_path)
            print(f"✅ 任务完成: {task_id}\n")
            return {
                "success": True,
                "task_id": task_id,
                "message": "处理完成",
                "download_url": f"/api/download/{task_id}"
            }
        else:
            print(f"❌ 任务失败: {message}\n")
            raise HTTPException(500, message)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, f"错误: {str(e)}")


@app.get("/api/download/{task_id}")
async def download(task_id: str):
    """下载结果"""
    if task_id not in tasks:
        raise HTTPException(404, "任务不存在")

    file_path = tasks[task_id]
    if not Path(file_path).exists():
        raise HTTPException(404, "文件不存在")

    return FileResponse(file_path, media_type='video/mp4', filename=f"output_{task_id}.mp4")


@app.get("/api/models")
async def models():
    """可用模型"""
    return {
        "models": ["inswapper_128", "ghost_1_256", "simswap_256", "hyperswap_1a_256"],
        "recommended": "inswapper_128"
    }


@app.get("/api/health")
async def health():
    """健康检查"""
    return {"status": "ok", "tasks": len(tasks)}


@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "FaceFusion API",
        "version": "1.0.0",
        "docs": "/docs"
    }


if __name__ == "__main__":
    print("=" * 60)
    print("🎬 FaceFusion FastAPI Server")
    print("=" * 60)

    # 初始化配置
    print("⚙️  初始化配置...")
    init_facefusion()
    _initialized = True

    print(f"📁 上传目录: {UPLOAD_DIR.absolute()}")
    print(f"📁 输出目录: {OUTPUT_DIR.absolute()}")
    print("=" * 60)
    print("🚀 启动中...")
    print("📚 API 文档: http://127.0.0.1:8000/docs")
    print("=" * 60)

    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
